# 训练模型
import base64
import sys
import time
import numpy as np

from keras.callbacks import ModelCheckpoint, EarlyStopping
# from tensorflow.keras.optimizers import SGD
from keras.optimizers import SGD
from sklearn.metrics import roc_curve, auc, precision_score, recall_score, f1_score, confusion_matrix, \
    precision_recall_curve
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
from ACRmodel.my_model import MyModel
from pathlib import Path
import os
import io
import zipfile
import requests  # 引入 requests 库
from utils.upload_image import upload_image_to_imgbed
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 0=全部信息, 1=忽略通知, 2=忽略警告, 3=忽略错误


def train_model(plant_name, dataset_id, model_id, epochs):
    """通用模型训练函数
        Args:
            plant_name: 植物种类名称
            dataset_id: 数据集id (数据集文件夹名字)
            model_id: 模型id (模型文件夹名称)
            epochs: 训练轮数
        Returns:
            模型的状态：未训练,训练中,训练失败,公开,私密
            模型的准确率
            模型的AUC
            模型的precision
            模型的recall
            模型的f1
            模型的指标图
            模型的混淆矩阵图
            模型的文件大小
        """
    train_result = {
        "status": "训练失败",
        "message": "服务器错误，模型训练失败",
        "accuracy": 0,
        "auc": 0,
        "precision": 0,
        "recall": 0,
        "f1": 0,
        "metrics_image": "",
        "confusion_matrix_image": "",
        "file_size": 0,
        "stop_epoch": 0
    }
    try:
        # 初始化随机种子和matplotlib配置
        np.random.seed(666)
        matplotlib.use('Agg')

        # 获取项目根目录
        current_file_path = os.path.abspath(__file__)  # 预期：F:\Project\PlantACR\ACRmodel\train.py
        model_root = os.path.dirname(os.path.dirname(current_file_path))
        print(f"项目根目录：{model_root}")  # 预期：F:\Project\PlantACR

        # 加载数据集
        print("\033[1;30;34m%s\033[0m" % "加载数据集中...")
        # 立即输出缓冲区
        sys.stdout.flush()
        try:
            data_dir = Path(model_root) / "ACRmodel" / "static" / "dataset" / plant_name / str(dataset_id)
            data_train = np.load(data_dir / "data_train.npy")
            label_train = np.load(data_dir / "label_train.npy")
            data_val = np.load(data_dir / "data_val.npy")
            label_val = np.load(data_dir / "label_val.npy")
            data_test = np.load(data_dir / "data_test.npy")
            label_test = np.load(data_dir / "label_test.npy")
            print("\033[1;30;34m%s\033[0m" % "加载数据完成！")
        except Exception as e:
            print("\033[1;30;34m%s\033[0m" % "加载数据失败：" + str(e))
            train_result["status"] = "训练失败"
            train_result["message"] = "加载数据集失败！"
            return train_result

        # 构件自定义模型
        model = MyModel()
        sgd = SGD(learning_rate=float(0.001), decay=1e-5, momentum=0.9, nesterov=True)
        model.compile(loss='binary_crossentropy', optimizer=sgd, metrics=['accuracy'])

        # 模型保存配置
        save_dir = Path(model_root) / "ACRmodel" / "static" / "weights" / plant_name / str(model_id)
        os.makedirs(save_dir, exist_ok=True)
        filepath = Path(save_dir) / "model_weights_best"
        sys.stdout.flush()  # 训练是长时间运行的任务，需要及时将运行状态输出给用户

        # 回调函数，用于在训练过程中实现早停（Early Stopping）机制
        # 在验证集连续20轮次损失值不再下降时停止训练，以避免过拟合
        early_stopping = EarlyStopping(monitor='val_loss', patience=20, verbose=1)

        # 回调函数，保存模型
        checkpointer = ModelCheckpoint(filepath, verbose=1, save_best_only=True, save_weights_only=True)

        # 模型训练
        print("\033[1;30;34m%s\033[0m" % "编译模型完成")
        print("\033[1;30;34m%s\033[0m" % "模型训练中...")
        time_start = time.time()
        result = model.fit(data_train, label_train, batch_size=64, epochs=epochs, shuffle=True,
                           validation_data=(data_val, label_val), callbacks=[checkpointer, early_stopping])
        time_end = time.time()

        # 获取实际训练轮次
        if early_stopping.stopped_epoch == 0:  # 没有触发早停
            stopped_epochs = epochs  # 使用设置的完整轮次
        else:  # 触发早停
            stopped_epochs = early_stopping.stopped_epoch + 1  # 因为索引从0开始
        print(f"\033[1;30;34m训练提前停止于第 {stopped_epochs} 轮\033[0m")
        print("\033[1;30;34m%s\033[0m" % "模型训练完成！")
        print('training time : %d sec' % (time_end - time_start))

        # 模型评估
        print("\033[1;30;34m%s\033[0m" % "评估模型：")
        model.load_weights(filepath)
        score_test = model.evaluate(data_test, label_test, verbose=0)
        print('测试集准确率 :', score_test[1])
        acc = score_test[1]

        # 后续评估和绘图
        pred = model.predict(data_test)
        # 阈值=0.5，计算二分类预测结果
        pred_classes = (pred > 0.5).astype("int32")

        # 计算指标
        precision = precision_score(label_test, pred_classes)
        recall = recall_score(label_test, pred_classes)
        f1 = f1_score(label_test, pred_classes)
        print(f"精确率(Precision): {precision:.4f}")
        print(f"召回率(Recall): {recall:.4f}")
        print(f"F1分数: {f1:.4f}")
        fpr, tpr, thresholds = roc_curve(label_test, pred)
        roc_auc = auc(fpr, tpr)  # 使用 auc 函数计算 ROC 曲线下的面积，这个值越接近 1，说明模型的性能越好
        print("ROC 曲线下的面积（AUC）:", roc_auc)

        # 绘制混淆矩阵
        print("\033[1;30;34m%s\033[0m" % "绘制混淆矩阵...")
        cm = confusion_matrix(label_test, pred_classes)
        plt.figure(figsize=(6, 6))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
        plt.title('Confusion Matrix')
        plt.ylabel('True label')
        plt.xlabel('Predicted label')
        plt.savefig(save_dir / 'confusion_matrix.png')
        plt.close()

        # 绘制 ROC 曲线，曲线越靠近左上角, AUC值越接近 1，模型的性能就越好
        print("\033[1;30;34m%s\033[0m" % "绘制ROC曲线图像...")
        plt.title('Receiver Operating Characteristic')
        # 蓝色曲线表示 ROC 曲线，红色虚线表示随机猜测的曲线
        plt.plot(fpr, tpr, 'b', label='AUC = %0.2f' % roc_auc)
        plt.legend(loc='lower right')
        plt.plot([0, 1], [0, 1], 'r--')
        plt.xlim([0, 1])
        plt.ylim([0, 1])
        # 图像的横轴是假正率，纵轴是真正率
        plt.ylabel('True Positive Rate')
        plt.xlabel('False Positive Rate')
        plt.savefig(save_dir / 'auc.png')
        plt.close()

        # 绘制模型的损失函数图像，并保存为图片文件
        print("\033[1;30;34m%s\033[0m" % "绘制损失函数图像...")
        plt.plot(result.history['loss'])
        plt.plot(result.history['val_loss'])
        plt.title('Model Loss')
        # 横轴是 epoch 数量，纵轴是损失函数值
        plt.ylabel('loss')
        plt.xlabel('epoch')
        plt.legend(['train', 'val'], loc='upper right')
        plt.savefig(save_dir / 'loss.png')
        plt.close()

        # 绘制模型的准确率图像，并保存为图片文件
        print("\033[1;30;34m%s\033[0m" % "绘制准确率图像...")
        plt.figure()
        # 横轴是 epoch 数量，纵轴是准确率
        # 蓝色曲线表示训练准确率，橙色曲线表示验证准确率
        plt.plot(result.epoch, result.history['accuracy'], label="accuracy")
        plt.plot(result.epoch, result.history['val_accuracy'], label="val_accuracy")
        plt.title('Model Accuracy')
        plt.scatter(result.epoch, result.history['accuracy'], marker='*')
        plt.scatter(result.epoch, result.history['val_accuracy'])
        plt.legend(loc='lower right')
        plt.savefig(save_dir / 'accuracy.png')
        plt.close()

        # 绘制精确率-召回率曲线
        print("\033[1;30;34m%s\033[0m" % "绘制PR曲线...")
        precision_curve, recall_curve, _ = precision_recall_curve(label_test, pred)
        plt.figure()
        plt.plot(recall_curve, precision_curve, label='PR Curve')
        plt.xlabel('Recall')
        plt.ylabel('Precision')
        plt.title('Precision-Recall Curve')
        plt.legend()
        plt.savefig(save_dir / 'pr_curve.png')
        plt.close()

        # 绘制合并图，创建2x2的子图布局
        plt.figure(figsize=(16, 12))
        # 1. 准确率图表 (左上)
        plt.subplot(2, 2, 1)
        plt.plot(result.epoch, result.history['accuracy'], label="accuracy")
        plt.plot(result.epoch, result.history['val_accuracy'], label="val_accuracy")
        plt.title('Model Accuracy')
        plt.scatter(result.epoch, result.history['accuracy'], marker='*')
        plt.scatter(result.epoch, result.history['val_accuracy'])
        plt.legend(loc='lower right')
        # 2. 损失函数图表 (右上)
        plt.subplot(2, 2, 2)
        plt.plot(result.history['loss'])
        plt.plot(result.history['val_loss'])
        plt.title('Model Loss')
        plt.ylabel('loss')
        plt.xlabel('epoch')
        plt.legend(['train', 'val'], loc='upper right')
        # 3. ROC曲线 (左下)
        plt.subplot(2, 2, 3)
        plt.title('Receiver Operating Characteristic')
        plt.plot(fpr, tpr, 'b', label='AUC = %0.2f' % roc_auc)
        plt.legend(loc='lower right')
        plt.plot([0, 1], [0, 1], 'r--')
        plt.xlim([0, 1])
        plt.ylim([0, 1])
        plt.ylabel('True Positive Rate')
        plt.xlabel('False Positive Rate')
        # 4. PR曲线 (右下)
        plt.subplot(2, 2, 4)
        precision_curve, recall_curve, _ = precision_recall_curve(label_test, pred)
        plt.plot(recall_curve, precision_curve, label='PR Curve')
        plt.xlabel('Recall')
        plt.ylabel('Precision')
        plt.title('Precision-Recall Curve')
        plt.legend()
        # 保存合并后的图表
        plt.tight_layout()
        plt.savefig(save_dir / 'metrics.png')
        plt.close()

        print("\033[1;30;34m%s\033[0m" % "所有图像绘制完成！")

        # 训练完成后释放GPU内存
        # K.clear_session()
        # tf.compat.v1.reset_default_graph()
        # gc.collect()

        # 计算模型文件总大小
        total_size = sum(os.path.getsize(os.path.join(save_dir, f)) for f in os.listdir(save_dir)
                         if os.path.isfile(os.path.join(save_dir, f)))
        file_size = f"{total_size / (1024 * 1024):.2f}MB"
        print(f"模型文件总大小: {file_size}")

        # 将所有模型文件打包成ZIP

        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for file in os.listdir(save_dir):
                file_path = os.path.join(save_dir, file)
                if os.path.isfile(file_path):  # 只打包文件，不打包子目录
                    zipf.write(file_path, arcname=file)  # 保持文件名不变

        # 保存ZIP文件
        zip_filename = f"model_{model_id}.zip"
        zip_path = os.path.join(save_dir, zip_filename)
        with open(zip_path, 'wb') as f:
            f.write(zip_buffer.getvalue())
        print(f"已生成模型压缩包: {zip_path}")

        # 读取图片并转换为base64
        def image_to_base64(image_path):
            with open(image_path, "rb") as image_file:
                return base64.b64encode(image_file.read()).decode('utf-8')

        try:
            # 上传图片到图床并获取URL
            metrics_image_url = upload_image_to_imgbed(save_dir / 'metrics.png')
            confusion_matrix_image_url = upload_image_to_imgbed(save_dir / 'confusion_matrix.png')
        except Exception as e:
            print("上传图片到图床失败：", e)
            metrics_image_url = ""
            confusion_matrix_image_url = ""

        # 返回评价指标
        train_result["status"] = "训练成功"
        train_result["message"] = "模型训练完成！"
        train_result["accuracy"] = acc
        train_result["auc"] = roc_auc
        train_result["precision"] = precision
        train_result["recall"] = recall
        train_result["f1"] = f1
        train_result["metrics_image"] = metrics_image_url
        train_result["confusion_matrix_image"] = confusion_matrix_image_url
        train_result["file_size"] = file_size
        train_result["stopped_epochs"] = stopped_epochs

        # 释放GPU内存
        # K.clear_session()
        # tf.compat.v1.reset_default_graph()
        # gc.collect()
        return train_result
    except Exception as e:
        print("训练模型时出现错误：", e)
        train_result["status"] = "训练失败"
        train_result["message"] = "训练模型时出现错误!"

        # 出错时也释放GPU内存
        # K.clear_session()
        # tf.compat.v1.reset_default_graph()
        # gc.collect()
        return train_result


if __name__ == "__main__":
    # 测试代码
    results = train_model("拟南芥", "1", "7", 2)
    print(results)
